import torch
from models import build_model
from main import get_args_parser
import argparse

if __name__=="__main__":
    parser = argparse.ArgumentParser('DETR training and evaluation script', parents=[get_args_parser()])
    args = parser.parse_args()
    model, criterion, postprocessors = build_model(args)
    input = torch.randn(1, 3, 800, 1200)
    out = model(input)
    print("pred_logits: ", out['pred_logits'].shape), # [1,100,92]
    print("pred_boxes: ", out['pred_boxes'].shape), # [1,100, 4]
    print("out['aux_outputs']: ", len(out['aux_outputs']), out['aux_outputs'][0]['pred_logits'].shape, 
          out['aux_outputs'][0]['pred_boxes'].shape)
    